import glob
from skimage import io
import cv2
import numpy as np
from matplotlib import pyplot as plt
import json
from torchvision import transforms as t_trans
from albumentations import *
from PIL import Image
import torch
import os


class Fpv(object):
    def __init__(self,  root='/home/dsl/dataset/fpv1007_4pixel/',  is_train=True):
        np.random.seed(100)
        self.root = root
        self.is_train = is_train
        '''
        if self.is_train:
            annpath = os.path.join(self.root, 'ImageSets','Segmentation' , 'train.txt')
            with open(annpath, 'r') as fr:
                pairs = fr.read().splitlines()
        else:
            annpath = os.path.join(self.root, 'ImageSets', 'Segmentation', 'val.txt')
            with open(annpath, 'r') as fr:
                pairs = fr.read().splitlines()
        '''
        pairs = [p[0:-4] for p in os.listdir(os.path.join(self.root,'JPEGImages'))]
        np.random.shuffle(pairs)
        lx = len(pairs)
        if self.is_train:
            pairs = pairs[0: int(0.9*lx)]
        else:
            pairs = pairs[int(0.9 * lx):]
            '''
            with open('val.txt', 'w') as f:
                for p in pairs:
                    f.write(p+'\n')
                    f.flush()
            '''

        self.img_paths, self.lb_paths = [], []
        for pair in pairs:
            self.img_paths.append(os.path.join(self.root, 'JPEGImages',pair+'.jpg'))
            self.lb_paths.append(os.path.join(self.root, 'SegmentationClass',pair+'.png'))

        self.train_aug = Compose([
            Resize(height=720, width=1280, p=1),
            #RandomScale(scale_limit=(0.05, 0.3), p=0.5),
            RandomCrop(height=599, width=1056 ,p= 0.5),
            RandomCrop(height=544, width=960, p=0.5),
            #RandomBrightnessContrast(p=0.3),
            RGBShift(r_shift_limit=20, g_shift_limit=20, b_shift_limit=20, p=0.3),
            ImageCompression(quality_lower=80, p=0.2),
            #MotionBlur(p=0.2),
            #CLAHE(p=1.0),
            #RandomFog(fog_coef_lower=0.1, fog_coef_upper=0.4, p=1),
            #Resize(height=512,width=960, p=1),
            HorizontalFlip(p=0.5),
        ])
        '''
        self.test_aug = Compose([
            Resize(height=720, width=1280, p=1),
            #Resize(height=544, width=960, p=1),
        ])
        '''
        self.trans = t_trans.Compose([
            t_trans.ToTensor(),
            #t_trans.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

    def __len__(self):
        return len(self.img_paths)

    def __getitem__(self, item):
        if True:
            img = cv2.imread(self.img_paths[item])[:, :, ::-1].copy()
            mask = cv2.imread(self.lb_paths[item], 0)

            if self.is_train:
                auged = self.train_aug(image=img, mask=mask)
                img = auged['image']
                mask = auged['mask']
            '''
            plt.subplot(121)
            plt.imshow(img)
            plt.subplot(122)
            plt.imshow(mask)
            plt.show()
            '''
            mask = mask[:, :]
            mask[mask == 4] = 0
            mask[mask == 8] = 0
            mask[mask == 9] = 4
            mask[mask == 10] = 5
            mask[mask == 11] = 6
            mask[mask == 14] = 7
            mask[mask == 15] = 8
            img = cv2.resize(img,dsize=(960, 544), interpolation=cv2.INTER_LINEAR)
            mask = cv2.resize(mask, dsize=(960, 544), interpolation=cv2.INTER_NEAREST)

            ig_data = Image.fromarray(img)
            ig_data = self.trans(ig_data)
            mask = torch.from_numpy(mask).long()
            return ig_data, mask


if __name__ == '__main__':
    dataset = Fpv(root='F:/1219', is_train=True)
    for i in range(100):
        dataset.__getitem__(i)



